import wave
import pyaudio
import pylab as pl
import numpy as np


class ProcessAudio:
    def __init__(self):
        self.pause = False
        pass

    def readwav(self, filehandle):
        params = filehandle.getparams()
        nchannels, sampwidth, samplerate, nsamples = params[:4]
        # read the data
        str_data = filehandle.readframes(nsamples)

        wave_data = np.fromstring(str_data, dtype=np.short)

        wave_data = wave_data * 1.0 / max(abs(wave_data))
        if params[0] == 2:
            wave_data.shape = -1, 2
            wave_data = wave_data.T
        new_wave = []
        # print(wave_data)
        for i, w in enumerate(wave_data):
            if len(w) > 3000000:
                new_wave.append((wave_data[i])[:3000000])
            else:
                new_wave.append(wave_data[i])
        return new_wave

    def writewav(self, outfilename, writemode, data, fs, nchannel):
        fileHandle = wave.open(outfilename, writemode)
        fileHandle.setnchannels(nchannel)
        fileHandle.setsampwidth(2)
        fileHandle.setframerate(fs)
        fileHandle.writeframes(data.tostring())
        fileHandle.close()

    def recordwav(self, filename, time, fs, nchannel):
        pa = pyaudio.PyAudio()
        save_buffer = ''
        buffer_size = 1000
        stream = pa.open(
            format=pyaudio.paInt16,
            channels=1,
            rate=fs,
            input=True,
            frames_per_buffer=buffer_size)

        read_time_per_second = fs / buffer_size
        cnt = 0
        while cnt < time * read_time_per_second:
            str_data = stream.read(buffer_size)
            save_buffer += str_data
            cnt += 1

        wave_data = np.fromstring(save_buffer, dtype=np.short)
        self.writewav(filename, "wb", wave_data, fs, nchannel)

    def playwav(self, filehandle):
        p = pyaudio.PyAudio()
        stream = p.open(
            format=p.get_format_from_width(filehandle.getsampwidth()),
            channels=filehandle.getnchannels(),
            rate=filehandle.getframerate(),
            output=True)
        data = filehandle.readframes(1024)
        while data != '':
            if self.pause:
                break
            stream.write(data)
            data = filehandle.readframes(1024)
        stream.close()
        p.terminate()

    def plotwav(self, data, fs, nchannel):
        if 2 == nchannel:
            length = len(data[0])
            time = np.arange(0, length) * (1.0 / fs)
            pl.subplot(211)
            pl.plot(time, data[0])
            pl.subplot(212)
            pl.plot(time, data[1], c="g")
            pl.xlabel("time (seconds)")
            pl.show()
        else:
            length = len(data)
            time = np.arange(0, length) * (1.0 / fs)
            pl.plot(time, data)
            pl.show()

    def plotspect(self, filehandle):
        params = filehandle.getparams()
        nchannels, sampwidth, framerate, nframes = params[:4]

        waveData = self.readwav(filehandle)

        framelength = 0.025  # 帧长20~30ms
        framesize = framelength * framerate  # 每帧点数 N = t*fs,通常情况下值为256或512,要与NFFT相等\
        # 而NFFT最好取2的整数次方,即framesize最好取的整数次方

        # 找到与当前framesize最接近的2的正整数次方
        nfftdict = {}
        lists = [32, 64, 128, 256, 512, 1024]
        for i in lists:
            nfftdict[i] = abs(framesize - i)
        sortlist = sorted(nfftdict.items(), key=lambda x: x[1])  # 按与当前framesize差值升序排列
        framesize = int(sortlist[0][0])  # 取最接近当前framesize的那个2的正整数次方值为新的framesize

        NFFT = framesize  # NFFT必须与时域的点数framsize相等，即不补零的FFT
        overlapSize = 1.0 / 3 * framesize  # 重叠部分采样点数overlapSize约为每帧点数的1/3~1/2
        overlapSize = int(round(overlapSize))  # 取整
        print("帧长为{},帧叠为{},傅里叶变换点数为{}".format(framesize, overlapSize, NFFT))
        spectrum, freqs, ts, fig = pl.specgram(waveData[0], NFFT=NFFT, Fs=framerate, window=np.hanning(M=framesize),
                                                noverlap=overlapSize, mode='default', scale_by_freq=True,
                                                sides='default', scale='dB', xextent=None)  # 绘制频谱图

        pl.ylabel('Frequency')
        pl.xlabel('Time')
        pl.title("Spectrogram")
        pl.show()

    def plot_spect(self, filehandle):
        params = filehandle.getparams()
        nchannels, sampwidth, framerate, nframes = params[:4]
        # 读取波形数据
        waveData = self.readwav(filehandle)

        N = 44100
        print(len(waveData[0]))
        start = int(len(waveData[0]) / 2)
        # 开始采样位置
        df = framerate / (N - 1)
        # 分辨率
        freq = [df * n for n in range(0, N)]
        # N个元素
        wave_data2 = waveData[0][start:start + N]
        c = np.fft.fft(wave_data2) * 2 / N
        # 常规显示采样频率一半的频谱
        pl.plot(freq[:round(len(freq) / 2)], abs(c[:round(len(c) / 2)]), 'r')
        pl.title('Freq')
        pl.xlabel("Freq/Hz")
        pl.show()

